import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np
import time


# 快速排序算法
def quick_sort(arr, low, high, steps):
    if low < high:
        pi = partition(arr, low, high, steps)

        quick_sort(arr, low, pi - 1, steps)
        quick_sort(arr, pi + 1, high, steps)


def partition(arr, low, high, steps):
    pivot = arr[high]
    i = low - 1
    for j in range(low, high):
        steps.append((arr.copy(), j, high))  # 记录当前比较的位置
        if arr[j] < pivot:
            # steps.append((arr.copy(), i, j))
            i += 1
            arr[i], arr[j] = arr[j], arr[i]
            steps.append((arr.copy(), i, j))  # 记录交换后的位置
    steps.append((arr.copy(), i + 1, high))
    arr[i + 1], arr[high] = arr[high], arr[i + 1]
    steps.append((arr.copy(), i + 1, high))  # 记录交换后的位置
    return i + 1


# 随机生成一个数组
arr = np.random.randint(0, 100, 20)
steps = []

# 获取排序步骤
quick_sort(arr.copy(), 0, len(arr) - 1, steps)

# 创建动画
fig, ax = plt.subplots()
ax.set_xlim(-1, len(arr))
ax.set_ylim(0, max(arr) + 10)

bar_rects = ax.bar(range(len(arr)), arr, align="edge")


def animate(i):
    current_step = steps[i][0]
    compare_idx1 = steps[i][1]
    compare_idx2 = steps[i][2]

    for idx, (rect, val) in enumerate(zip(bar_rects, current_step)):
        if idx == compare_idx1 or idx == compare_idx2:
            time.sleep(0.05)
            rect.set_color('r')  # 将比较的柱子设为红色
        else:
            rect.set_color('b')  # 其余柱子设为蓝色
        rect.set_height(val)

    return bar_rects


ani = animation.FuncAnimation(fig, animate, frames=len(steps), interval=50, blit=True)

plt.show()
